#!/usr/bin/env python3
"""修复 diffusers transformer_z_image.py 的 NPU complex64 索引问题"""
import os
import sys

try:
    import diffusers
    target = os.path.join(os.path.dirname(diffusers.__file__), 
                         'models/transformers/transformer_z_image.py')
    
    if not os.path.exists(target):
        print(f"File not found: {target}")
        sys.exit(0)
    
    print(f"Reading {target}...")
    with open(target, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if already patched
    if '_npu_patched_v3' in content:
        print("Already patched (v3)")
        sys.exit(0)
    
    lines_list = content.split('\n')
    
    # Find the problematic line directly (more robust)
    target_line_idx = None
    for i, line in enumerate(lines_list):
        if 'result.append(self.freqs_cis[i][index])' in line:
            target_line_idx = i
            print(f"Found target line at {i+1}: {line.strip()}")
            break
    
    if target_line_idx is None:
        print("ERROR: Target line 'result.append(self.freqs_cis[i][index])' not found")
        print("Searching for similar patterns...")
        for i, line in enumerate(lines_list):
            if 'freqs_cis' in line and 'index' in line and 'result' in line:
                print(f"  Line {i+1}: {line.strip()}")
        sys.exit(1)
    
    # Find the loop that contains this line
    loop_start_idx = None
    for i in range(target_line_idx, max(0, target_line_idx - 50), -1):
        if 'for i in range' in lines_list[i]:
            loop_start_idx = i
            print(f"Found loop at line {i+1}: {lines_list[i].strip()}")
            break
    
    if loop_start_idx is None:
        print("WARNING: Could not find loop, will add patch before target line")
        loop_start_idx = target_line_idx
    
    # Get indentation
    target_line = lines_list[target_line_idx]
    indent = len(target_line) - len(target_line.lstrip())
    indent_str = ' ' * indent
    
    loop_line = lines_list[loop_start_idx]
    loop_indent = len(loop_line) - len(loop_line.lstrip())
    loop_indent_str = ' ' * loop_indent
    
    # Build new content
    new_lines = []
    conversion_added = False
    
    for i, line in enumerate(lines_list):
        # Add conversion logic before the loop
        if i == loop_start_idx and not conversion_added:
            new_lines.append(f'{loop_indent_str}# NPU patch v3: convert complex64 to real representation before loop')
            new_lines.append(f'{loop_indent_str}if not hasattr(self, "_freqs_cis_npu_converted"):')
            new_lines.append(f'{loop_indent_str}    if self.freqs_cis is not None and len(self.freqs_cis) > 0:')
            new_lines.append(f'{loop_indent_str}        first = self.freqs_cis[0]')
            new_lines.append(f'{loop_indent_str}        if hasattr(first, "dtype") and "complex" in str(first.dtype):')
            new_lines.append(f'{loop_indent_str}            real_freqs = []')
            new_lines.append(f'{loop_indent_str}            for fc in self.freqs_cis:')
            new_lines.append(f'{loop_indent_str}                fc_cpu = fc.cpu()')
            new_lines.append(f'{loop_indent_str}                real_freqs.append(torch.stack([fc_cpu.real, fc_cpu.imag], dim=-1))')
            new_lines.append(f'{loop_indent_str}            self.freqs_cis = real_freqs')
            new_lines.append(f'{loop_indent_str}    self._freqs_cis_npu_converted = True')
            new_lines.append('')
            conversion_added = True
        
        # Replace the problematic line
        if i == target_line_idx:
            new_lines.append(f'{indent_str}# NPU patch v3: index on CPU and reconstruct complex')
            new_lines.append(f'{indent_str}orig_dev = index.device')
            new_lines.append(f'{indent_str}idx_cpu = index.cpu()')
            new_lines.append(f'{indent_str}freq_data = self.freqs_cis[i][idx_cpu]')
            new_lines.append(f'{indent_str}if len(freq_data.shape) > 0 and freq_data.shape[-1] == 2:')
            new_lines.append(f'{indent_str}    freq_complex = torch.complex(freq_data[..., 0], freq_data[..., 1])')
            new_lines.append(f'{indent_str}else:')
            new_lines.append(f'{indent_str}    freq_complex = freq_data')
            new_lines.append(f'{indent_str}if orig_dev.type != "cpu":')
            new_lines.append(f'{indent_str}    freq_complex = freq_complex.to(orig_dev)')
            new_lines.append(f'{indent_str}result.append(freq_complex)')
            continue
        
        new_lines.append(line)
    
    # Write patched content
    new_content = '\n'.join(new_lines)
    
    with open(target, 'w', encoding='utf-8') as f:
        f.write(new_content)
    print(f"Successfully patched {target}")
    
    # Verify
    with open(target, 'r', encoding='utf-8') as f:
        verify = f.read()
    
    if 'result.append(self.freqs_cis[i][index])' in verify:
        print("ERROR: Original line still exists! Patch failed.")
        sys.exit(1)
    
    if '_freqs_cis_npu_converted' not in verify or 'orig_dev = index.device' not in verify:
        print("ERROR: Patch markers not found! Patch may have failed.")
        sys.exit(1)
    
    print("Patch verified successfully")
        
except Exception as e:
    import traceback
    print(f"Error: {e}", file=sys.stderr)
    traceback.print_exc(file=sys.stderr)
    sys.exit(1)
